import numpy as np
import pandas as pd
def solve_2(x0:float,ul:float,ur:float):
    '''
    求解分两片常数初值的黎曼问题（Burgers方程弱解）
    
    该函数根据给定的间断面位置和左右区域常数值 ，
    返回弱解函数及间断面信息。
    
    参数:
        x0 (float): 间断面初始位置
        ul (float): 左侧区域常数值
        ur (float): 右侧区域常数值
    
    返回:
        u (function): 时空二元函数
    '''
    if ul==ur:
        def u(x,t):
            x,t=np.broadcast_arrays(x,t)
            return np.full_like(x,ul)
        u.d=()
    elif ul>ur:
        s=(ul+ur)/2
        f=lambda t:x0+s*t
        f.d,f.domain=1,(0,np.inf)
        def u(x,t):
            x,t=np.broadcast_arrays(x,t)
            return np.where(x<f(t),ul,ur)
        u.d=(f,)
    elif ul<ur:
        g1,g2=lambda t:x0+ul*t,lambda t:x0+ur*t
        g1.d,g1.domain,g2.d,g2.domain=0,(0,np.inf),0,(0,np.inf)
        def u(x,t):
            x,t=np.broadcast_arrays(x,t)
            condlist=[x<g1(t),(g1(t)<=x)&(x<g2(t)),g2(t)<=x]
            choicelist=[ul,(x-x0)/t,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(g1,g2)
    u.argv=(x0,ul,ur)
    u.slice=0.5,1
    return u

def solve_3(xl:float,xr:float,ul:float,um:float,ur:float):
    '''
    求解分三片常数初值的黎曼问题（Burgers方程弱解）
    
    该函数根据给定的间断面位置和左右区域常数值 ，
    返回弱解函数及间断面信息。
    
    参数:
        xl,xr (float): 间断面初始位置
        ul,um,ur (float): 左、中、右区域常数值
    
    返回:
        u (function): 时空二元函数
    '''
    if xl==xr:
        return solve_2(xl,ul,ur)
    xl,xr=min(xl,xr),max(xl,xr)
    s,sl,sr=(ul+ur)/2,(ul+um)/2,(um+ur)/2
    if um<ul<=ur:# A
        t1=-2*(xr-xl)/(um-ul)
        x1=xl+sl*t1
        P=(x1-ul*t1-xr)/np.sqrt(t1)
        f1=lambda t:xl+sl*t# s
        f1.d,f1.domain=1,(0,t1)
        f2=lambda t:xr+um*t# w
        f2.d,f2.domain=0,(0,t1)
        f3=lambda t:xr+ur*t# w
        f3.d,f3.domain=0,(0,np.inf)
        f4=lambda t:xr+ul*t+P*np.sqrt(t)# w
        f4.d,f4.domain=0,(t1,np.inf)
        def u(x,t):
            condlist=[(t<t1)&(x<f1(t)),
                      (t<t1)&(f1(t)<=x)&(x<=f2(t)),
                      (t<t1)&(f2(t)<x)&(x<f3(t)),
                      (t<t1)&(f3(t)<=x),
                      (t>=t1)&(x<=f4(t)),
                      (t>=t1)&(f4(t)<x)&(x<f3(t)),
                      (t>=t1)&(f3(t)<=x)]
            choicelist=[ul,um,(x-xr)/t,ur,ul,(x-xr)/t,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(f1,f2,f3,f4)
        u.slice=t1/2,t1+1
    elif ul<=ur<um:# B
        t1=-2*(xr-xl)/(ur-um)
        x1=xr+sr*t1
        Q=(x1-ur*t1-xl)/np.sqrt(t1)
        g1=lambda t:xl+ul*t# w
        g1.d,g1.domain=0,(0,np.inf)
        g2=lambda t:xl+um*t# w
        g2.d,g2.domain=0,(0,t1)
        g3=lambda t:xr+sr*t# s
        g3.d,g3.domain=1,(0,t1)
        g4=lambda t:xl+ur*t+Q*np.sqrt(t)# w
        g4.d,g4.domain=0,(t1,np.inf)
        def u(x,t):
            condlist=[(t<t1)&(x<g1(t)),
                      (t<t1)&(g1(t)<x)&(x<g2(t)),
                      (t<t1)&(g2(t)<=x)&(x<=g3(t)),
                      (t<t1)&(g3(t)<x),
                      (t>=t1)&(x<=g1(t)),
                      (t>=t1)&(g1(t)<x)&(x<g4(t)),
                      (t>=t1)&(g4(t)<=x)]
            choicelist=[ul,(x-xl)/t,um,ur,ul,(x-xl)/t,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(g1,g2,g3,g4)
        u.slice=t1/2,t1+1
    elif um<ur<ul:# C
        t1=-2*(xr-xl)/(um-ul)
        x1=xl+sl*t1
        P=(x1-ul*t1-xr)/np.sqrt(t1)
        t2=P**2/(ul-ur)**2
        x2=xr+ur*t2
        f1=lambda t:xl+sl*t# s
        f1.d,f1.domain=1,(0,t1)
        f2=lambda t:xr+um*t# w
        f2.d,f2.domain=0,(0,t1)
        f3=lambda t:xr+ur*t# w
        f3.d,f3.domain=0,(0,t2)
        f4=lambda t:xr+ul*t+P*np.sqrt(t)# w
        f4.d,f4.domain=0,(t1,t2)
        f5=lambda t:x2+s*(t-t2)# s
        f5.d,f5.domain=1,(t2,np.inf)
        def u(x,t):
            condlist=[(t<t1)&(x<f1(t)),
                      (t<t1)&(f1(t)<=x)&(x<=f2(t)),
                      (t<t1)&(f2(t)<x)&(x<f3(t)),
                      (t<t1)&(f3(t)<=x),
                      (t>=t1)&(t<=t2)&(x<=f4(t)),
                      (t>=t1)&(t<=t2)&(f4(t)<x)&(x<f3(t)),
                      (t>=t1)&(t<=t2)&(f3(t)<=x),
                      (t>t2)&(x<f5(t)),
                      (t>t2)&(x>=f5(t))]
            choicelist=[ul,um,(x-xr)/t,ur,ul,(x-xr)/t,ur,ul,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(f1,f2,f3,f4,f5)
        u.slice=t1/2,(t1+t2)/2,t2+1
    elif ur<ul<um:# D
        t1=-2*(xr-xl)/(ur-um)
        x1=xr+sr*t1
        Q=(x1-ur*t1-xl)/np.sqrt(t1)
        t2=Q**2/(ul-ur)**2
        x2=xl+ul*t2
        g1=lambda t:xl+ul*t# w
        g1.d,g1.domain=0,(0,t2)
        g2=lambda t:xl+um*t# w
        g2.d,g2.domain=0,(0,t1)
        g3=lambda t:xr+sr*t# s
        g3.d,g3.domain=1,(0,t1)
        g4=lambda t:xl+ur*t+Q*np.sqrt(t)# w
        g4.d,g4.domain=0,(t1,t2)
        g5=lambda t:x2+s*(t-t2)# s
        g5.d,g5.domain=1,(t2,np.inf)
        def u(x,t):
            condlist=[(t<t1)&(x<=g1(t)),
                      (t<t1)&(g1(t)<x)&(x<g2(t)),
                      (t<t1)&(g2(t)<=x)&(x<=g3(t)),
                      (t<t1)&(g3(t)<x),
                      (t>=t1)&(t<=t2)&(x<=g1(t)),
                      (t>=t1)&(t<=t2)&(g1(t)<x)&(x<g4(t)),
                      (t>=t1)&(t<=t2)&(g4(t)<=x),
                      (t>t2)&(x<g5(t)),
                      (t>t2)&(x>=g5(t))]
            choicelist=[ul,(x-xl)/t,um,ur,ul,(x-xl)/t,ur,ul,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(g1,g2,g3,g4,g5)
        u.slice=t1/2,(t1+t2)/2,t2+1
    elif ur<um<ul:# E
        t1=-(xr-xl)/(sr-sl)
        x1=xl+sl*t1
        f1=lambda t:xl+sl*t# s
        f1.d,f1.domain=1,(0,t1)
        f2=lambda t:xr+sr*t# s
        f2.d,f2.domain=1,(0,t1)
        f3=lambda t:x1+s*(t-t1)# s
        f3.d,f3.domain=1,(t1,np.inf)
        def u(x,t):
            condlist=[(t<t1)&(x<f1(t)),
                      (t<t1)&(f1(t)<x)&(x<f2(t)),
                      (t<t1)&(f2(t)<x),
                      (t>=t1)&(x<f3(t)),
                      (t>=t1)&(f3(t)<=x)]
            choicelist=[ul,um,ur,ul,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(f1,f2,f3)
        u.slice=t1/2,t1+1
    elif ul<um<ur:# F
        t1=-(xr-xl)/(sr-sl)
        x1=xr+sr*t1
        g1=lambda t:xl+ul*t# w
        g1.d,g1.domain=0,(0,np.inf)
        g2=lambda t:xl+um*t# w
        g2.d,g2.domain=0,(0,np.inf)
        g3=lambda t:xr+um*t# w
        g3.d,g3.domain=0,(0,np.inf)
        g4=lambda t:xr+ur*t# W
        g4.d,g4.domain=0,(0,np.inf)
        def u(x,t):
            condlist=[x<g1(t),
                      (g1(t)<x)&(x<g2(t)),
                      (g2(t)<x)&(x<g3(t)),
                      (g3(t)<x)&(x<g4(t)),
                      g4(t)<x]
            choicelist=[ul,(x-xl)/t,um,(x-xr)/t,ur]
            return np.select(condlist,choicelist,np.nan)
        u.d=(g1,g2,g3,g4)
        u.slice=2.5,7.5
    elif ul==um:
        return solve_2(xr,ul,ur)
    elif um==ur:
        return solve_2(xl,ul,ur)
    else:
        return '未考虑情况'
    u.argv=(xl,xr,ul,um,ur)
    return u

def calculate(u:tuple,points:pd.DataFrame)->pd.DataFrame:
    '''
    根据给定的黎曼问题解函数，计算指定坐标点的物理量值
    
    参数:
        u (function): 时空二元函数 u(x, t)，返回对应点的物理量值
        
        points (pd.DataFrame): 包含计算坐标的 DataFrame，需包含列:
            - 'x': 空间坐标
            - 't': 时间坐标
    
    返回:
        pd.DataFrame: 输入 DataFrame 的副本，新增 'u' 列存储计算结果
            - 'u': 对应 (x, t) 坐标的物理量值，通过 u(x, t) 计算
    '''
    new_point=points.copy()
    new_point['u']=u(new_point['x'].values,new_point['t'].values)
    return new_point